def columnar_encrypt(plaintext, key):
    """
    列换位密码加密
    :param plaintext: 明文字符串
    :param key: 密钥字符串（如"CAT"）
    :return: 密文
    """
    plaintext = plaintext.replace(" ", "").upper()
    key = key.upper()
    key_length = len(key)

    # 生成列顺序（按密钥字母升序排列对应的列索引） 
    sorted_key_chars = sorted(key)
    key_order = [key.index(c) for c in sorted_key_chars]

    # 计算填充长度 [[10]]
    padding_length = (key_length - (len(plaintext) % key_length)) % key_length
    plaintext += 'X' * padding_length

    # 将明文按行填充到矩阵 
    rows = [plaintext[i:i + key_length] for i in range(0, len(plaintext), key_length)]

    # 按列顺序读取密文 
    ciphertext = []
    for col in key_order:
        for row in rows:
            ciphertext.append(row[col])
    return ''.join(ciphertext)


def columnar_decrypt(ciphertext, key):
    """
    列换位密码解密
    :param ciphertext: 密文字符串
    :param key: 加密时使用的密钥字符串
    :return: 明文（可能含填充的X）
    """
    ciphertext = ciphertext.upper()
    key = key.upper()
    key_length = len(key)
    cols = len(ciphertext) // key_length  # 每列字符数

    # 生成逆置换顺序 
    sorted_key_chars = sorted(key)
    key_order = [key.index(c) for c in sorted_key_chars]
    # 创建逆置换索引（加密列顺序 → 原始列顺序）
    inverse_order = [0] * key_length
    for new_idx, orig_idx in enumerate(key_order):
        inverse_order[orig_idx] = new_idx

    # 将密文分割为列 
    columns = [ciphertext[i * cols:(i + 1) * cols] for i in range(key_length)]

    # 按逆顺序重组列 
    ordered_cols = [columns[inverse_order[i]] for i in range(key_length)]

    # 按行读取恢复明文
    plaintext = []
    for i in range(cols):
        for col in ordered_cols:
            plaintext.append(col[i])
    return ''.join(plaintext).rstrip('X')  # 移除填充的X


# 使用示例 
if __name__ == "__main__":
    key = "CAT"  # 密钥字母顺序为A(1), C(0), T(2)
    message = "HELLOWORLD"

    encrypted = columnar_encrypt(message, key)
    print(f"加密结果: {encrypted}")  # 输出：EORXHLODLWLX

    decrypted = columnar_decrypt(encrypted, key)
    print(f"解密结果: {decrypted}")  # 输出：HELLOWORLD